{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import copy\n",
    "\n",
    "import numpy\n",
    "from IPython.display import display\n",
    "\n",
    "from rmgpy.tools.uncertainty import Uncertainty, ThermoParameterUncertainty, KineticParameterUncertainty"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Parameter Uncertainty Assignment\n",
    "\n",
    "This script aims to extract model sources in a clear and informative format.  The script first shows what all the kinetic and thermo sources are in a model. Then it goes through each reaction and species to show their source and what the assigned uncertainties are.  This can be used with any RMG-generated CHEMKIN file that is annotated.\n",
    "\n",
    "__Note__: The RMG-database version must match the version used to generate the model. RMG will attempt to recreate the kinetics estimate for each reaction and may fail if the database is different."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "chem_file = './data/parse_source/chem_annotated.inp'\n",
    "dict_file = './data/parse_source/species_dictionary.txt'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "uncertainty = Uncertainty(output_directory='./temp/uncertainty')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "uncertainty.load_model(chem_file, dict_file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# NOTE: You must load the database with the same settings which were used to generate the model.\n",
    "#       This includes any thermo or kinetics libraries which were used.\n",
    "uncertainty.load_database(\n",
    "    thermo_libraries=['DFT_QCI_thermo', 'primaryThermoLibrary'],\n",
    "    kinetics_families='default',\n",
    "    reaction_libraries=[],\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "uncertainty.extract_sources_from_model()\n",
    "uncertainty.compile_all_sources()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "print('All Kinetic Sources')\n",
    "for sourceType in uncertainty.all_kinetic_sources.keys():\n",
    "    if sourceType == 'Library':\n",
    "        print('============')\n",
    "        print('Library kinetics')\n",
    "        print('')\n",
    "        print('\\tReactions: ', uncertainty.all_kinetic_sources['Library'])\n",
    "    elif sourceType == 'PDep':\n",
    "        print('============')\n",
    "        print('PDep kinetics')\n",
    "        print('')\n",
    "        print('\\tReactions: ', uncertainty.all_kinetic_sources['PDep'])\n",
    "    elif sourceType == 'Rate Rules':\n",
    "        print('============')\n",
    "        print('Rate rule kinetics')\n",
    "        print('')\n",
    "        for familyLabel, entries in uncertainty.all_kinetic_sources['Rate Rules'].items():\n",
    "            print('\\t', familyLabel)\n",
    "            for entry in entries:\n",
    "                print('\\t\\t', entry)\n",
    "    elif sourceType == 'Training':\n",
    "        print('============')\n",
    "        print('Training reaction kinetics')\n",
    "        print('')\n",
    "        for familyLabel, entries in uncertainty.all_kinetic_sources['Training'].items():\n",
    "            print('\\t', familyLabel)\n",
    "            for entry in entries:\n",
    "                print('\\t\\t', entry)\n",
    "    else:\n",
    "        print(sourceType)\n",
    "        raise Exception('Kinetics source must be Library, PDep, Rate Rules, or Training')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('All Thermo Sources')\n",
    "for sourceType in uncertainty.all_thermo_sources.keys():\n",
    "    if sourceType == 'Library':\n",
    "        print('============')\n",
    "        print('Library thermo')\n",
    "        print('')\n",
    "        print('\\tSpecies: ', uncertainty.all_thermo_sources['Library'])\n",
    "    elif sourceType == 'QM':\n",
    "        print('============')\n",
    "        print('QM thermo')\n",
    "        print('')\n",
    "        print('\\tSpecies: ', uncertainty.all_thermo_sources['QM'])\n",
    "    elif sourceType == 'GAV':\n",
    "        print('============')\n",
    "        print('Group additivity thermo')\n",
    "        print('')\n",
    "        for groupType, entries in uncertainty.all_thermo_sources['GAV'].items():\n",
    "            print('\\t', groupType)\n",
    "            for entry in entries:\n",
    "                print('\\t\\t', entry)\n",
    "    else:\n",
    "        raise Exception('Thermo source must be GAV, QM, or Library')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Assign all the uncertainties using the Uncertainty() class function\n",
    "# ThermoParameterUncertainty and KineticParameterUncertainty classes may be customized and passed into this function\n",
    "# if non-default constants for constructing the uncertainties are desired\n",
    "uncertainty.assign_parameter_uncertainties()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "T = 623 # temperature in Kelvin for which to evaluate kinetics\n",
    "P = 1e5  # Pa "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "g_param_engine = ThermoParameterUncertainty()\n",
    "k_param_engine = KineticParameterUncertainty()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for rxn, source in uncertainty.reaction_sources_dict.items():\n",
    "    print('======')\n",
    "    print(rxn)\n",
    "    display(rxn)\n",
    "    if 'Library' in source:\n",
    "        print('Library reaction')\n",
    "        print(source['Library'])\n",
    "    elif 'PDep' in source:\n",
    "        print('PDep reaction')\n",
    "        print(source['PDep'])\n",
    "    elif 'Rate Rules' in source:\n",
    "        print('Rate rule estimate')\n",
    "        family = source['Rate Rules'][0]\n",
    "        sourceDict = source['Rate Rules'][1]\n",
    "        originalTemplate = sourceDict['template']\n",
    "        print('\\tFamily = ', family)\n",
    "        print('\\tOriginal Template = ', [group.label for group in originalTemplate])\n",
    "        print('\\tExact = ', sourceDict['exact'])\n",
    "        rules = sourceDict['rules']\n",
    "        training = sourceDict['training']\n",
    "        if rules:\n",
    "            print('\\tRate rule sources:')\n",
    "            for ruleEntry, weight in rules:\n",
    "                print('\\t\\t', ruleEntry, '=', weight)\n",
    "        if training:\n",
    "            print('\\tTraining sources:')\n",
    "            for ruleEntry, trainingEntry, weight in training:\n",
    "                print('\\t\\t', ruleEntry , 'mapped to', trainingEntry , '=', weight)\n",
    "    elif 'Training' in source:\n",
    "        print('Training reaction')\n",
    "        family = source['Training'][0]\n",
    "        training = source['Training'][1]\n",
    "        print('\\t Family = ', family)\n",
    "        print('\\t\\t', training)\n",
    "\n",
    "    print('')\n",
    "    print('Rate coefficient at {} K = {:.2e}'.format(T, rxn.kinetics.get_rate_coefficient(T,P)))\n",
    "\n",
    "    # Uncomment the following lines if you want to verify that the parsing has been performed correctly by\n",
    "    # checking the values for both the original and reconstructed kinetics\n",
    "#     print('---------')\n",
    "#     print('Original kinetics:')\n",
    "#     print(rxn.kinetics)\n",
    "#     print('')\n",
    "#     print('Reconstructed kinetics from parsing:')\n",
    "#     reconstructedKinetics = uncertainty.database.kinetics.reconstruct_kinetics_from_source(rxn, source, fix_barrier_height=True)\n",
    "#     print(reconstructedKinetics)\n",
    "\n",
    "#     rxnIndex = uncertainty.reaction_list.index(rxn)\n",
    "#     print('Uncertainty dln(k) = ', uncertainty.kinetic_input_uncertainties[rxnIndex])\n",
    "    \n",
    "#     # Test that the partial uncertainty calculation is working\n",
    "#     dlnk = 0.0\n",
    "#     if 'Rate Rules' in source:\n",
    "#         family = source['Rate Rules'][0]\n",
    "#         sourceDict = source['Rate Rules'][1]\n",
    "#         rules = sourceDict['rules']\n",
    "#         training = sourceDict['training']\n",
    "#         for ruleEntry, weight in rules:\n",
    "#             dlnk += k_param_engine.get_partial_uncertainty_value(source, 'Rate Rules', corr_param=ruleEntry, corr_family=family)\n",
    "#         for ruleEntry, trainingEntry, weight in training:\n",
    "#             dlnk += k_param_engine.get_partial_uncertainty_value(source, 'Rate Rules', corr_param=ruleEntry, corr_family=family)\n",
    "#         dlnk += k_param_engine.get_partial_uncertainty_value(source, 'Estimation')\n",
    "#     elif 'PDep' in source:\n",
    "#         dlnk += k_param_engine.get_partial_uncertainty_value(source, 'PDep', source['PDep'])\n",
    "#     elif 'Library' in source:\n",
    "#         dlnk += k_param_engine.get_partial_uncertainty_value(source, 'Library', source['Library'])\n",
    "#     elif 'Training' in source:\n",
    "#         dlnk += k_param_engine.get_partial_uncertainty_value(source, 'Training', source['Training'])\n",
    "#     print('Uncertainty dlnk calculated using sum of partial values = ', dlnk)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for species, source in uncertainty.species_sources_dict.items():\n",
    "    print('==========')\n",
    "    print(species)\n",
    "    display(species)\n",
    "    if 'Library' in source:\n",
    "        print('Thermo Library: ', source['Library'])\n",
    "    if 'QM' in source:\n",
    "        print('QM: ', source['QM'])\n",
    "    if 'GAV' in source:\n",
    "        print('Group additivity:')\n",
    "        for groupType, groupList in source['GAV'].items():\n",
    "            print('\\t', groupType)\n",
    "            for group, weight in groupList:\n",
    "                print('\\t\\t', group, '=', weight)\n",
    "\n",
    "    spcIndex = uncertainty.species_list.index(species)    \n",
    "    print('')\n",
    "    print('Uncertainty dG = ', uncertainty.thermo_input_uncertainties[spcIndex], ' kcal/mol')\n",
    "    \n",
    "    \n",
    "    # Test that the partial uncertainty calculation is working\n",
    "    dG = 0.0\n",
    "    if 'Library' in source:\n",
    "        dG += g_param_engine.get_partial_uncertainty_value(source, 'Library', corr_param=source['Library'])\n",
    "    if 'QM' in source:\n",
    "        dG += g_param_engine.get_partial_uncertainty_value(source, 'QM',corr_param=source['QM'])\n",
    "    if 'GAV' in source:\n",
    "        for groupType, groupList in source['GAV'].items():\n",
    "            for group, weight in groupList:\n",
    "                dG += g_param_engine.get_partial_uncertainty_value(source, 'GAV', group, groupType)\n",
    "        dG += g_param_engine.get_partial_uncertainty_value(source, 'Estimation')\n",
    "    print('Uncertainty dG calculated using sum of partial values = ', dG, ' kcal/mol')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Assign correlated parameter uncertainties \n",
    "uncertainty.assign_parameter_uncertainties(correlated=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# See the thermo correlated parameter partial uncertainties\n",
    "uncertainty.thermo_input_uncertainties"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "# See the kinetics correlated parameter partial uncertainties\n",
    "uncertainty.kinetic_input_uncertainties"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:rmg_env]",
   "language": "python",
   "name": "conda-env-rmg_env-py"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
